import torch
import numpy as np
from mmdet.datasets.builder import PIPELINES

@PIPELINES.register_module()
class ProcessNSTPGraph:
    """处理nSTP图数据的转换组件"""
    
    def __init__(self, graph_feat_dim=64, with_agent_type=True):
        self.graph_feat_dim = graph_feat_dim
        self.with_agent_type = with_agent_type
    
    def __call__(self, results):
        """处理nSTP图数据"""
        if 'nstp_graph' not in results:
            return results
            
        graph_data = results['nstp_graph']
        
        # 处理PyG Data对象
        if hasattr(graph_data, 'x') and hasattr(graph_data, 'edge_index'):
            # 已经是PyG Data对象，确保张量类型正确
            if not isinstance(graph_data.x, torch.Tensor):
                graph_data.x = torch.tensor(graph_data.x, dtype=torch.float)
            if not isinstance(graph_data.edge_index, torch.Tensor):
                graph_data.edge_index = torch.tensor(graph_data.edge_index, dtype=torch.long)
            if hasattr(graph_data, 'edge_attr') and not isinstance(graph_data.edge_attr, torch.Tensor):
                graph_data.edge_attr = torch.tensor(graph_data.edge_attr, dtype=torch.float)
                
        # 处理字典格式的图数据
        elif isinstance(graph_data, dict):
            # 处理节点特征
            if 'x' in graph_data:
                x = graph_data['x']
                if isinstance(x, np.ndarray):
                    x = torch.from_numpy(x).float()
                elif isinstance(x, list):
                    x = torch.tensor(x).float()
                elif not isinstance(x, torch.Tensor):
                    x = torch.tensor(x, dtype=torch.float)
                graph_data['x'] = x
                
            # 处理边索引
            if 'edge_index' in graph_data:
                edge_index = graph_data['edge_index']
                if isinstance(edge_index, np.ndarray):
                    edge_index = torch.from_numpy(edge_index).long()
                elif isinstance(edge_index, list):
                    edge_index = torch.tensor(edge_index).long()
                elif not isinstance(edge_index, torch.Tensor):
                    edge_index = torch.tensor(edge_index, dtype=torch.long)
                graph_data['edge_index'] = edge_index
                
            # 处理边属性
            if 'edge_attr' in graph_data:
                edge_attr = graph_data['edge_attr']
                if isinstance(edge_attr, np.ndarray):
                    edge_attr = torch.from_numpy(edge_attr).float()
                elif isinstance(edge_attr, list):
                    edge_attr = torch.tensor(edge_attr).float()
                elif not isinstance(edge_attr, torch.Tensor):
                    edge_attr = torch.tensor(edge_attr, dtype=torch.float)
                graph_data['edge_attr'] = edge_attr
        
        # 更新结果
        results['nstp_graph'] = graph_data
        return results